import os
import json
from typing import List, Dict
import pandas as pd


def process_jsonl(jsonl_path):
    pass


def calculate_hilight(content:str, summary:List[str]):
    hilight_info = []
    lines = content.split("\n")
    for i, line in enumerate(lines):
        for long_sen in summary:
            for short_sen in long_sen.split():
                if len(short_sen) < 3:
                    continue
                row = i + 1
                idx = line.find(short_sen)
                if idx == -1:
                    continue
                col_start = idx
                col_end = col_start + len(short_sen)
                info = {'row': row, 'col_start': col_start, 'col_end': col_end}
                hilight_info.append(info)
    return hilight_info


def process_csv(csv_path):
    df = pd.read_csv(csv_path)
    assert "query" in df.columns
    assert "summary" in df.columns
    assert "snippet" in df.columns
    assert "url" in df.columns
    assert "core_content" in df.columns
    if len(df) == 0:
        return []
    # 统一中英文字符
    char_map = {
        '？': '?',
        '，': ',',
        '。': '.',
        '：': ':',
        '“': '"',
        '”': '"',
        '‘': "'",
        '’': "'",
        '；': ';',
        '（': '(',
        '）': ')',
        '【': '[',
        '】': ']'
    }
    for zh, en in char_map.items():
        df["query"] = df["query"].apply(lambda x: x.replace(zh, en))
        df["core_content"] = df["core_content"].apply(lambda x: x.replace(zh, en))
    
    if isinstance(df["summary"][0], str):
        df["summary"] = df["summary"].apply(lambda x: json.loads(x))
    df["core_content"] = df["core_content"].apply(lambda x: x.replace("\\n", "\n"))
    meta_info_list = []
    for i, row in df.iterrows():
        if row['query'] == '2013年的长轴高顶依维柯价值多少':
            print("here")
        hilight_info = calculate_hilight(row["core_content"], row["summary"])
        meta_info = {
            "query": row["query"],
            "summary": row["summary"],
            "snippet": row["snippet"],
            "url": row["url"],
            "core_content": row["core_content"],
            "hilight_info": hilight_info
        }
        meta_info_list.append(meta_info)
    return meta_info_list



def test_process_csv():
    csv_path = os.path.join(os.path.dirname(__file__), "data", "sampled_data_xlmr_v2 - Sheet1.csv")
    meta_info_list = process_csv(csv_path)
    print(meta_info_list[:2])


def main():
    test_process_csv()


if __name__ == '__main__':
    main()